{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Help on GRU in module torch.nn.modules.rnn object:\n",
      "\n",
      "class GRU(RNNBase)\n",
      " |  GRU(*args, **kwargs)\n",
      " |  \n",
      " |  __init__(input_size,hidden_size,num_layers=1,bias=True,batch_first=False,dropout=0.0,bidirectional=False,device=None,dtype=None)\n",
      " |  \n",
      " |  Apply a multi-layer gated recurrent unit (GRU) RNN to an input sequence.\n",
      " |  For each element in the input sequence, each layer computes the following\n",
      " |  function:\n",
      " |  \n",
      " |  .. math::\n",
      " |      \\begin{array}{ll}\n",
      " |          r_t = \\sigma(W_{ir} x_t + b_{ir} + W_{hr} h_{(t-1)} + b_{hr}) \\\\\n",
      " |          z_t = \\sigma(W_{iz} x_t + b_{iz} + W_{hz} h_{(t-1)} + b_{hz}) \\\\\n",
      " |          n_t = \\tanh(W_{in} x_t + b_{in} + r_t \\odot (W_{hn} h_{(t-1)}+ b_{hn})) \\\\\n",
      " |          h_t = (1 - z_t) \\odot n_t + z_t \\odot h_{(t-1)}\n",
      " |      \\end{array}\n",
      " |  \n",
      " |  where :math:`h_t` is the hidden state at time `t`, :math:`x_t` is the input\n",
      " |  at time `t`, :math:`h_{(t-1)}` is the hidden state of the layer\n",
      " |  at time `t-1` or the initial hidden state at time `0`, and :math:`r_t`,\n",
      " |  :math:`z_t`, :math:`n_t` are the reset, update, and new gates, respectively.\n",
      " |  :math:`\\sigma` is the sigmoid function, and :math:`\\odot` is the Hadamard product.\n",
      " |  \n",
      " |  In a multilayer GRU, the input :math:`x^{(l)}_t` of the :math:`l` -th layer\n",
      " |  (:math:`l \\ge 2`) is the hidden state :math:`h^{(l-1)}_t` of the previous layer multiplied by\n",
      " |  dropout :math:`\\delta^{(l-1)}_t` where each :math:`\\delta^{(l-1)}_t` is a Bernoulli random\n",
      " |  variable which is :math:`0` with probability :attr:`dropout`.\n",
      " |  \n",
      " |  Args:\n",
      " |      input_size: The number of expected features in the input `x`\n",
      " |      hidden_size: The number of features in the hidden state `h`\n",
      " |      num_layers: Number of recurrent layers. E.g., setting ``num_layers=2``\n",
      " |          would mean stacking two GRUs together to form a `stacked GRU`,\n",
      " |          with the second GRU taking in outputs of the first GRU and\n",
      " |          computing the final results. Default: 1\n",
      " |      bias: If ``False``, then the layer does not use bias weights `b_ih` and `b_hh`.\n",
      " |          Default: ``True``\n",
      " |      batch_first: If ``True``, then the input and output tensors are provided\n",
      " |          as `(batch, seq, feature)` instead of `(seq, batch, feature)`.\n",
      " |          Note that this does not apply to hidden or cell states. See the\n",
      " |          Inputs/Outputs sections below for details.  Default: ``False``\n",
      " |      dropout: If non-zero, introduces a `Dropout` layer on the outputs of each\n",
      " |          GRU layer except the last layer, with dropout probability equal to\n",
      " |          :attr:`dropout`. Default: 0\n",
      " |      bidirectional: If ``True``, becomes a bidirectional GRU. Default: ``False``\n",
      " |  \n",
      " |  Inputs: input, h_0\n",
      " |      * **input**: tensor of shape :math:`(L, H_{in})` for unbatched input,\n",
      " |        :math:`(L, N, H_{in})` when ``batch_first=False`` or\n",
      " |        :math:`(N, L, H_{in})` when ``batch_first=True`` containing the features of\n",
      " |        the input sequence.  The input can also be a packed variable length sequence.\n",
      " |        See :func:`torch.nn.utils.rnn.pack_padded_sequence` or\n",
      " |        :func:`torch.nn.utils.rnn.pack_sequence` for details.\n",
      " |      * **h_0**: tensor of shape :math:`(D * \\text{num\\_layers}, H_{out})` or\n",
      " |        :math:`(D * \\text{num\\_layers}, N, H_{out})`\n",
      " |        containing the initial hidden state for the input sequence. Defaults to zeros if not provided.\n",
      " |  \n",
      " |      where:\n",
      " |  \n",
      " |      .. math::\n",
      " |          \\begin{aligned}\n",
      " |              N ={} & \\text{batch size} \\\\\n",
      " |              L ={} & \\text{sequence length} \\\\\n",
      " |              D ={} & 2 \\text{ if bidirectional=True otherwise } 1 \\\\\n",
      " |              H_{in} ={} & \\text{input\\_size} \\\\\n",
      " |              H_{out} ={} & \\text{hidden\\_size}\n",
      " |          \\end{aligned}\n",
      " |  \n",
      " |  Outputs: output, h_n\n",
      " |      * **output**: tensor of shape :math:`(L, D * H_{out})` for unbatched input,\n",
      " |        :math:`(L, N, D * H_{out})` when ``batch_first=False`` or\n",
      " |        :math:`(N, L, D * H_{out})` when ``batch_first=True`` containing the output features\n",
      " |        `(h_t)` from the last layer of the GRU, for each `t`. If a\n",
      " |        :class:`torch.nn.utils.rnn.PackedSequence` has been given as the input, the output\n",
      " |        will also be a packed sequence.\n",
      " |      * **h_n**: tensor of shape :math:`(D * \\text{num\\_layers}, H_{out})` or\n",
      " |        :math:`(D * \\text{num\\_layers}, N, H_{out})` containing the final hidden state\n",
      " |        for the input sequence.\n",
      " |  \n",
      " |  Attributes:\n",
      " |      weight_ih_l[k] : the learnable input-hidden weights of the :math:`\\text{k}^{th}` layer\n",
      " |          (W_ir|W_iz|W_in), of shape `(3*hidden_size, input_size)` for `k = 0`.\n",
      " |          Otherwise, the shape is `(3*hidden_size, num_directions * hidden_size)`\n",
      " |      weight_hh_l[k] : the learnable hidden-hidden weights of the :math:`\\text{k}^{th}` layer\n",
      " |          (W_hr|W_hz|W_hn), of shape `(3*hidden_size, hidden_size)`\n",
      " |      bias_ih_l[k] : the learnable input-hidden bias of the :math:`\\text{k}^{th}` layer\n",
      " |          (b_ir|b_iz|b_in), of shape `(3*hidden_size)`\n",
      " |      bias_hh_l[k] : the learnable hidden-hidden bias of the :math:`\\text{k}^{th}` layer\n",
      " |          (b_hr|b_hz|b_hn), of shape `(3*hidden_size)`\n",
      " |  \n",
      " |  .. note::\n",
      " |      All the weights and biases are initialized from :math:`\\mathcal{U}(-\\sqrt{k}, \\sqrt{k})`\n",
      " |      where :math:`k = \\frac{1}{\\text{hidden\\_size}}`\n",
      " |  \n",
      " |  .. note::\n",
      " |      For bidirectional GRUs, forward and backward are directions 0 and 1 respectively.\n",
      " |      Example of splitting the output layers when ``batch_first=False``:\n",
      " |      ``output.view(seq_len, batch, num_directions, hidden_size)``.\n",
      " |  \n",
      " |  .. note::\n",
      " |      ``batch_first`` argument is ignored for unbatched inputs.\n",
      " |  \n",
      " |  .. note::\n",
      " |      The calculation of new gate :math:`n_t` subtly differs from the original paper and other frameworks.\n",
      " |      In the original implementation, the Hadamard product :math:`(\\odot)` between :math:`r_t` and the\n",
      " |      previous hidden state :math:`h_{(t-1)}` is done before the multiplication with the weight matrix\n",
      " |      `W` and addition of bias:\n",
      " |  \n",
      " |      .. math::\n",
      " |          \\begin{aligned}\n",
      " |              n_t = \\tanh(W_{in} x_t + b_{in} + W_{hn} ( r_t \\odot h_{(t-1)} ) + b_{hn})\n",
      " |          \\end{aligned}\n",
      " |  \n",
      " |      This is in contrast to PyTorch implementation, which is done after :math:`W_{hn} h_{(t-1)}`\n",
      " |  \n",
      " |      .. math::\n",
      " |          \\begin{aligned}\n",
      " |              n_t = \\tanh(W_{in} x_t + b_{in} + r_t \\odot (W_{hn} h_{(t-1)}+ b_{hn}))\n",
      " |          \\end{aligned}\n",
      " |  \n",
      " |      This implementation differs on purpose for efficiency.\n",
      " |  \n",
      " |  .. include:: ../cudnn_persistent_rnn.rst\n",
      " |  \n",
      " |  Examples::\n",
      " |  \n",
      " |      >>> rnn = nn.GRU(10, 20, 2)\n",
      " |      >>> input = torch.randn(5, 3, 10)\n",
      " |      >>> h0 = torch.randn(2, 3, 20)\n",
      " |      >>> output, hn = rnn(input, h0)\n",
      " |  \n",
      " |  Method resolution order:\n",
      " |      GRU\n",
      " |      RNNBase\n",
      " |      torch.nn.modules.module.Module\n",
      " |      builtins.object\n",
      " |  \n",
      " |  Methods defined here:\n",
      " |  \n",
      " |  __init__(self, *args, **kwargs)\n",
      " |      Initialize internal Module state, shared by both nn.Module and ScriptModule.\n",
      " |  \n",
      " |  forward(self, input, hx=None)\n",
      " |      Define the computation performed at every call.\n",
      " |      \n",
      " |      Should be overridden by all subclasses.\n",
      " |      \n",
      " |      .. note::\n",
      " |          Although the recipe for forward pass needs to be defined within\n",
      " |          this function, one should call the :class:`Module` instance afterwards\n",
      " |          instead of this since the former takes care of running the\n",
      " |          registered hooks while the latter silently ignores them.\n",
      " |  \n",
      " |  ----------------------------------------------------------------------\n",
      " |  Data and other attributes defined here:\n",
      " |  \n",
      " |  __annotations__ = {}\n",
      " |  \n",
      " |  ----------------------------------------------------------------------\n",
      " |  Methods inherited from RNNBase:\n",
      " |  \n",
      " |  __getstate__(self)\n",
      " |      Helper for pickle.\n",
      " |  \n",
      " |  __setattr__(self, attr, value)\n",
      " |      Implement setattr(self, name, value).\n",
      " |  \n",
      " |  __setstate__(self, d)\n",
      " |  \n",
      " |  check_forward_args(self, input: torch.Tensor, hidden: torch.Tensor, batch_sizes: Optional[torch.Tensor])\n",
      " |  \n",
      " |  check_hidden_size(self, hx: torch.Tensor, expected_hidden_size: Tuple[int, int, int], msg: str = 'Expected hidden size {}, got {}') -> None\n",
      " |  \n",
      " |  check_input(self, input: torch.Tensor, batch_sizes: Optional[torch.Tensor]) -> None\n",
      " |  \n",
      " |  extra_repr(self) -> str\n",
      " |      Return the extra representation of the module.\n",
      " |      \n",
      " |      To print customized extra information, you should re-implement\n",
      " |      this method in your own modules. Both single-line and multi-line\n",
      " |      strings are acceptable.\n",
      " |  \n",
      " |  flatten_parameters(self) -> None\n",
      " |      Reset parameter data pointer so that they can use faster code paths.\n",
      " |      \n",
      " |      Right now, this works only if the module is on the GPU and cuDNN is enabled.\n",
      " |      Otherwise, it's a no-op.\n",
      " |  \n",
      " |  get_expected_hidden_size(self, input: torch.Tensor, batch_sizes: Optional[torch.Tensor]) -> Tuple[int, int, int]\n",
      " |  \n",
      " |  permute_hidden(self, hx: torch.Tensor, permutation: Optional[torch.Tensor])\n",
      " |  \n",
      " |  reset_parameters(self) -> None\n",
      " |  \n",
      " |  ----------------------------------------------------------------------\n",
      " |  Readonly properties inherited from RNNBase:\n",
      " |  \n",
      " |  all_weights\n",
      " |  \n",
      " |  ----------------------------------------------------------------------\n",
      " |  Data and other attributes inherited from RNNBase:\n",
      " |  \n",
      " |  __constants__ = ['mode', 'input_size', 'hidden_size', 'num_layers', 'b...\n",
      " |  \n",
      " |  __jit_unused_properties__ = ['all_weights']\n",
      " |  \n",
      " |  ----------------------------------------------------------------------\n",
      " |  Methods inherited from torch.nn.modules.module.Module:\n",
      " |  \n",
      " |  __call__ = _wrapped_call_impl(self, *args, **kwargs)\n",
      " |  \n",
      " |  __delattr__(self, name)\n",
      " |      Implement delattr(self, name).\n",
      " |  \n",
      " |  __dir__(self)\n",
      " |      Default dir() implementation.\n",
      " |  \n",
      " |  __getattr__(self, name: str) -> Union[torch.Tensor, ForwardRef('Module')]\n",
      " |      # It is crucial that the return type is not annotated as `Any`, otherwise type checking\n",
      " |      # on `torch.nn.Module` and all its subclasses is largely disabled as a result. See:\n",
      " |      # https://github.com/pytorch/pytorch/pull/115074\n",
      " |  \n",
      " |  __repr__(self)\n",
      " |      Return repr(self).\n",
      " |  \n",
      " |  add_module(self, name: str, module: Optional[ForwardRef('Module')]) -> None\n",
      " |      Add a child module to the current module.\n",
      " |      \n",
      " |      The module can be accessed as an attribute using the given name.\n",
      " |      \n",
      " |      Args:\n",
      " |          name (str): name of the child module. The child module can be\n",
      " |              accessed from this module using the given name\n",
      " |          module (Module): child module to be added to the module.\n",
      " |  \n",
      " |  apply(self: ~T, fn: Callable[[ForwardRef('Module')], NoneType]) -> ~T\n",
      " |      Apply ``fn`` recursively to every submodule (as returned by ``.children()``) as well as self.\n",
      " |      \n",
      " |      Typical use includes initializing the parameters of a model\n",
      " |      (see also :ref:`nn-init-doc`).\n",
      " |      \n",
      " |      Args:\n",
      " |          fn (:class:`Module` -> None): function to be applied to each submodule\n",
      " |      \n",
      " |      Returns:\n",
      " |          Module: self\n",
      " |      \n",
      " |      Example::\n",
      " |      \n",
      " |          >>> @torch.no_grad()\n",
      " |          >>> def init_weights(m):\n",
      " |          >>>     print(m)\n",
      " |          >>>     if type(m) == nn.Linear:\n",
      " |          >>>         m.weight.fill_(1.0)\n",
      " |          >>>         print(m.weight)\n",
      " |          >>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))\n",
      " |          >>> net.apply(init_weights)\n",
      " |          Linear(in_features=2, out_features=2, bias=True)\n",
      " |          Parameter containing:\n",
      " |          tensor([[1., 1.],\n",
      " |                  [1., 1.]], requires_grad=True)\n",
      " |          Linear(in_features=2, out_features=2, bias=True)\n",
      " |          Parameter containing:\n",
      " |          tensor([[1., 1.],\n",
      " |                  [1., 1.]], requires_grad=True)\n",
      " |          Sequential(\n",
      " |            (0): Linear(in_features=2, out_features=2, bias=True)\n",
      " |            (1): Linear(in_features=2, out_features=2, bias=True)\n",
      " |          )\n",
      " |  \n",
      " |  bfloat16(self: ~T) -> ~T\n",
      " |      Casts all floating point parameters and buffers to ``bfloat16`` datatype.\n",
      " |      \n",
      " |      .. note::\n",
      " |          This method modifies the module in-place.\n",
      " |      \n",
      " |      Returns:\n",
      " |          Module: self\n",
      " |  \n",
      " |  buffers(self, recurse: bool = True) -> Iterator[torch.Tensor]\n",
      " |      Return an iterator over module buffers.\n",
      " |      \n",
      " |      Args:\n",
      " |          recurse (bool): if True, then yields buffers of this module\n",
      " |              and all submodules. Otherwise, yields only buffers that\n",
      " |              are direct members of this module.\n",
      " |      \n",
      " |      Yields:\n",
      " |          torch.Tensor: module buffer\n",
      " |      \n",
      " |      Example::\n",
      " |      \n",
      " |          >>> # xdoctest: +SKIP(\"undefined vars\")\n",
      " |          >>> for buf in model.buffers():\n",
      " |          >>>     print(type(buf), buf.size())\n",
      " |          <class 'torch.Tensor'> (20L,)\n",
      " |          <class 'torch.Tensor'> (20L, 1L, 5L, 5L)\n",
      " |  \n",
      " |  children(self) -> Iterator[ForwardRef('Module')]\n",
      " |      Return an iterator over immediate children modules.\n",
      " |      \n",
      " |      Yields:\n",
      " |          Module: a child module\n",
      " |  \n",
      " |  compile(self, *args, **kwargs)\n",
      " |      Compile this Module's forward using :func:`torch.compile`.\n",
      " |      \n",
      " |      This Module's `__call__` method is compiled and all arguments are passed as-is\n",
      " |      to :func:`torch.compile`.\n",
      " |      \n",
      " |      See :func:`torch.compile` for details on the arguments for this function.\n",
      " |  \n",
      " |  cpu(self: ~T) -> ~T\n",
      " |      Move all model parameters and buffers to the CPU.\n",
      " |      \n",
      " |      .. note::\n",
      " |          This method modifies the module in-place.\n",
      " |      \n",
      " |      Returns:\n",
      " |          Module: self\n",
      " |  \n",
      " |  cuda(self: ~T, device: Union[int, torch.device, NoneType] = None) -> ~T\n",
      " |      Move all model parameters and buffers to the GPU.\n",
      " |      \n",
      " |      This also makes associated parameters and buffers different objects. So\n",
      " |      it should be called before constructing the optimizer if the module will\n",
      " |      live on GPU while being optimized.\n",
      " |      \n",
      " |      .. note::\n",
      " |          This method modifies the module in-place.\n",
      " |      \n",
      " |      Args:\n",
      " |          device (int, optional): if specified, all parameters will be\n",
      " |              copied to that device\n",
      " |      \n",
      " |      Returns:\n",
      " |          Module: self\n",
      " |  \n",
      " |  double(self: ~T) -> ~T\n",
      " |      Casts all floating point parameters and buffers to ``double`` datatype.\n",
      " |      \n",
      " |      .. note::\n",
      " |          This method modifies the module in-place.\n",
      " |      \n",
      " |      Returns:\n",
      " |          Module: self\n",
      " |  \n",
      " |  eval(self: ~T) -> ~T\n",
      " |      Set the module in evaluation mode.\n",
      " |      \n",
      " |      This has an effect only on certain modules. See the documentation of\n",
      " |      particular modules for details of their behaviors in training/evaluation\n",
      " |      mode, i.e. whether they are affected, e.g. :class:`Dropout`, :class:`BatchNorm`,\n",
      " |      etc.\n",
      " |      \n",
      " |      This is equivalent with :meth:`self.train(False) <torch.nn.Module.train>`.\n",
      " |      \n",
      " |      See :ref:`locally-disable-grad-doc` for a comparison between\n",
      " |      `.eval()` and several similar mechanisms that may be confused with it.\n",
      " |      \n",
      " |      Returns:\n",
      " |          Module: self\n",
      " |  \n",
      " |  float(self: ~T) -> ~T\n",
      " |      Casts all floating point parameters and buffers to ``float`` datatype.\n",
      " |      \n",
      " |      .. note::\n",
      " |          This method modifies the module in-place.\n",
      " |      \n",
      " |      Returns:\n",
      " |          Module: self\n",
      " |  \n",
      " |  get_buffer(self, target: str) -> 'Tensor'\n",
      " |      Return the buffer given by ``target`` if it exists, otherwise throw an error.\n",
      " |      \n",
      " |      See the docstring for ``get_submodule`` for a more detailed\n",
      " |      explanation of this method's functionality as well as how to\n",
      " |      correctly specify ``target``.\n",
      " |      \n",
      " |      Args:\n",
      " |          target: The fully-qualified string name of the buffer\n",
      " |              to look for. (See ``get_submodule`` for how to specify a\n",
      " |              fully-qualified string.)\n",
      " |      \n",
      " |      Returns:\n",
      " |          torch.Tensor: The buffer referenced by ``target``\n",
      " |      \n",
      " |      Raises:\n",
      " |          AttributeError: If the target string references an invalid\n",
      " |              path or resolves to something that is not a\n",
      " |              buffer\n",
      " |  \n",
      " |  get_extra_state(self) -> Any\n",
      " |      Return any extra state to include in the module's state_dict.\n",
      " |      \n",
      " |      Implement this and a corresponding :func:`set_extra_state` for your module\n",
      " |      if you need to store extra state. This function is called when building the\n",
      " |      module's `state_dict()`.\n",
      " |      \n",
      " |      Note that extra state should be picklable to ensure working serialization\n",
      " |      of the state_dict. We only provide backwards compatibility guarantees\n",
      " |      for serializing Tensors; other objects may break backwards compatibility if\n",
      " |      their serialized pickled form changes.\n",
      " |      \n",
      " |      Returns:\n",
      " |          object: Any extra state to store in the module's state_dict\n",
      " |  \n",
      " |  get_parameter(self, target: str) -> 'Parameter'\n",
      " |      Return the parameter given by ``target`` if it exists, otherwise throw an error.\n",
      " |      \n",
      " |      See the docstring for ``get_submodule`` for a more detailed\n",
      " |      explanation of this method's functionality as well as how to\n",
      " |      correctly specify ``target``.\n",
      " |      \n",
      " |      Args:\n",
      " |          target: The fully-qualified string name of the Parameter\n",
      " |              to look for. (See ``get_submodule`` for how to specify a\n",
      " |              fully-qualified string.)\n",
      " |      \n",
      " |      Returns:\n",
      " |          torch.nn.Parameter: The Parameter referenced by ``target``\n",
      " |      \n",
      " |      Raises:\n",
      " |          AttributeError: If the target string references an invalid\n",
      " |              path or resolves to something that is not an\n",
      " |              ``nn.Parameter``\n",
      " |  \n",
      " |  get_submodule(self, target: str) -> 'Module'\n",
      " |      Return the submodule given by ``target`` if it exists, otherwise throw an error.\n",
      " |      \n",
      " |      For example, let's say you have an ``nn.Module`` ``A`` that\n",
      " |      looks like this:\n",
      " |      \n",
      " |      .. code-block:: text\n",
      " |      \n",
      " |          A(\n",
      " |              (net_b): Module(\n",
      " |                  (net_c): Module(\n",
      " |                      (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))\n",
      " |                  )\n",
      " |                  (linear): Linear(in_features=100, out_features=200, bias=True)\n",
      " |              )\n",
      " |          )\n",
      " |      \n",
      " |      (The diagram shows an ``nn.Module`` ``A``. ``A`` which has a nested\n",
      " |      submodule ``net_b``, which itself has two submodules ``net_c``\n",
      " |      and ``linear``. ``net_c`` then has a submodule ``conv``.)\n",
      " |      \n",
      " |      To check whether or not we have the ``linear`` submodule, we\n",
      " |      would call ``get_submodule(\"net_b.linear\")``. To check whether\n",
      " |      we have the ``conv`` submodule, we would call\n",
      " |      ``get_submodule(\"net_b.net_c.conv\")``.\n",
      " |      \n",
      " |      The runtime of ``get_submodule`` is bounded by the degree\n",
      " |      of module nesting in ``target``. A query against\n",
      " |      ``named_modules`` achieves the same result, but it is O(N) in\n",
      " |      the number of transitive modules. So, for a simple check to see\n",
      " |      if some submodule exists, ``get_submodule`` should always be\n",
      " |      used.\n",
      " |      \n",
      " |      Args:\n",
      " |          target: The fully-qualified string name of the submodule\n",
      " |              to look for. (See above example for how to specify a\n",
      " |              fully-qualified string.)\n",
      " |      \n",
      " |      Returns:\n",
      " |          torch.nn.Module: The submodule referenced by ``target``\n",
      " |      \n",
      " |      Raises:\n",
      " |          AttributeError: If the target string references an invalid\n",
      " |              path or resolves to something that is not an\n",
      " |              ``nn.Module``\n",
      " |  \n",
      " |  half(self: ~T) -> ~T\n",
      " |      Casts all floating point parameters and buffers to ``half`` datatype.\n",
      " |      \n",
      " |      .. note::\n",
      " |          This method modifies the module in-place.\n",
      " |      \n",
      " |      Returns:\n",
      " |          Module: self\n",
      " |  \n",
      " |  ipu(self: ~T, device: Union[int, torch.device, NoneType] = None) -> ~T\n",
      " |      Move all model parameters and buffers to the IPU.\n",
      " |      \n",
      " |      This also makes associated parameters and buffers different objects. So\n",
      " |      it should be called before constructing the optimizer if the module will\n",
      " |      live on IPU while being optimized.\n",
      " |      \n",
      " |      .. note::\n",
      " |          This method modifies the module in-place.\n",
      " |      \n",
      " |      Arguments:\n",
      " |          device (int, optional): if specified, all parameters will be\n",
      " |              copied to that device\n",
      " |      \n",
      " |      Returns:\n",
      " |          Module: self\n",
      " |  \n",
      " |  load_state_dict(self, state_dict: Mapping[str, Any], strict: bool = True, assign: bool = False)\n",
      " |      Copy parameters and buffers from :attr:`state_dict` into this module and its descendants.\n",
      " |      \n",
      " |      If :attr:`strict` is ``True``, then\n",
      " |      the keys of :attr:`state_dict` must exactly match the keys returned\n",
      " |      by this module's :meth:`~torch.nn.Module.state_dict` function.\n",
      " |      \n",
      " |      .. warning::\n",
      " |          If :attr:`assign` is ``True`` the optimizer must be created after\n",
      " |          the call to :attr:`load_state_dict` unless\n",
      " |          :func:`~torch.__future__.get_swap_module_params_on_conversion` is ``True``.\n",
      " |      \n",
      " |      Args:\n",
      " |          state_dict (dict): a dict containing parameters and\n",
      " |              persistent buffers.\n",
      " |          strict (bool, optional): whether to strictly enforce that the keys\n",
      " |              in :attr:`state_dict` match the keys returned by this module's\n",
      " |              :meth:`~torch.nn.Module.state_dict` function. Default: ``True``\n",
      " |          assign (bool, optional): When set to ``False``, the properties of the tensors\n",
      " |              in the current module are preserved whereas setting it to ``True`` preserves\n",
      " |              properties of the Tensors in the state dict. The only\n",
      " |              exception is the ``requires_grad`` field of :class:`~torch.nn.Parameter`s\n",
      " |              for which the value from the module is preserved.\n",
      " |              Default: ``False``\n",
      " |      \n",
      " |      Returns:\n",
      " |          ``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields:\n",
      " |              * **missing_keys** is a list of str containing any keys that are expected\n",
      " |                  by this module but missing from the provided ``state_dict``.\n",
      " |              * **unexpected_keys** is a list of str containing the keys that are not\n",
      " |                  expected by this module but present in the provided ``state_dict``.\n",
      " |      \n",
      " |      Note:\n",
      " |          If a parameter or buffer is registered as ``None`` and its corresponding key\n",
      " |          exists in :attr:`state_dict`, :meth:`load_state_dict` will raise a\n",
      " |          ``RuntimeError``.\n",
      " |  \n",
      " |  modules(self) -> Iterator[ForwardRef('Module')]\n",
      " |      Return an iterator over all modules in the network.\n",
      " |      \n",
      " |      Yields:\n",
      " |          Module: a module in the network\n",
      " |      \n",
      " |      Note:\n",
      " |          Duplicate modules are returned only once. In the following\n",
      " |          example, ``l`` will be returned only once.\n",
      " |      \n",
      " |      Example::\n",
      " |      \n",
      " |          >>> l = nn.Linear(2, 2)\n",
      " |          >>> net = nn.Sequential(l, l)\n",
      " |          >>> for idx, m in enumerate(net.modules()):\n",
      " |          ...     print(idx, '->', m)\n",
      " |      \n",
      " |          0 -> Sequential(\n",
      " |            (0): Linear(in_features=2, out_features=2, bias=True)\n",
      " |            (1): Linear(in_features=2, out_features=2, bias=True)\n",
      " |          )\n",
      " |          1 -> Linear(in_features=2, out_features=2, bias=True)\n",
      " |  \n",
      " |  mtia(self: ~T, device: Union[int, torch.device, NoneType] = None) -> ~T\n",
      " |      Move all model parameters and buffers to the MTIA.\n",
      " |      \n",
      " |      This also makes associated parameters and buffers different objects. So\n",
      " |      it should be called before constructing the optimizer if the module will\n",
      " |      live on MTIA while being optimized.\n",
      " |      \n",
      " |      .. note::\n",
      " |          This method modifies the module in-place.\n",
      " |      \n",
      " |      Arguments:\n",
      " |          device (int, optional): if specified, all parameters will be\n",
      " |              copied to that device\n",
      " |      \n",
      " |      Returns:\n",
      " |          Module: self\n",
      " |  \n",
      " |  named_buffers(self, prefix: str = '', recurse: bool = True, remove_duplicate: bool = True) -> Iterator[Tuple[str, torch.Tensor]]\n",
      " |      Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.\n",
      " |      \n",
      " |      Args:\n",
      " |          prefix (str): prefix to prepend to all buffer names.\n",
      " |          recurse (bool, optional): if True, then yields buffers of this module\n",
      " |              and all submodules. Otherwise, yields only buffers that\n",
      " |              are direct members of this module. Defaults to True.\n",
      " |          remove_duplicate (bool, optional): whether to remove the duplicated buffers in the result. Defaults to True.\n",
      " |      \n",
      " |      Yields:\n",
      " |          (str, torch.Tensor): Tuple containing the name and buffer\n",
      " |      \n",
      " |      Example::\n",
      " |      \n",
      " |          >>> # xdoctest: +SKIP(\"undefined vars\")\n",
      " |          >>> for name, buf in self.named_buffers():\n",
      " |          >>>     if name in ['running_var']:\n",
      " |          >>>         print(buf.size())\n",
      " |  \n",
      " |  named_children(self) -> Iterator[Tuple[str, ForwardRef('Module')]]\n",
      " |      Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.\n",
      " |      \n",
      " |      Yields:\n",
      " |          (str, Module): Tuple containing a name and child module\n",
      " |      \n",
      " |      Example::\n",
      " |      \n",
      " |          >>> # xdoctest: +SKIP(\"undefined vars\")\n",
      " |          >>> for name, module in model.named_children():\n",
      " |          >>>     if name in ['conv4', 'conv5']:\n",
      " |          >>>         print(module)\n",
      " |  \n",
      " |  named_modules(self, memo: Optional[Set[ForwardRef('Module')]] = None, prefix: str = '', remove_duplicate: bool = True)\n",
      " |      Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.\n",
      " |      \n",
      " |      Args:\n",
      " |          memo: a memo to store the set of modules already added to the result\n",
      " |          prefix: a prefix that will be added to the name of the module\n",
      " |          remove_duplicate: whether to remove the duplicated module instances in the result\n",
      " |              or not\n",
      " |      \n",
      " |      Yields:\n",
      " |          (str, Module): Tuple of name and module\n",
      " |      \n",
      " |      Note:\n",
      " |          Duplicate modules are returned only once. In the following\n",
      " |          example, ``l`` will be returned only once.\n",
      " |      \n",
      " |      Example::\n",
      " |      \n",
      " |          >>> l = nn.Linear(2, 2)\n",
      " |          >>> net = nn.Sequential(l, l)\n",
      " |          >>> for idx, m in enumerate(net.named_modules()):\n",
      " |          ...     print(idx, '->', m)\n",
      " |      \n",
      " |          0 -> ('', Sequential(\n",
      " |            (0): Linear(in_features=2, out_features=2, bias=True)\n",
      " |            (1): Linear(in_features=2, out_features=2, bias=True)\n",
      " |          ))\n",
      " |          1 -> ('0', Linear(in_features=2, out_features=2, bias=True))\n",
      " |  \n",
      " |  named_parameters(self, prefix: str = '', recurse: bool = True, remove_duplicate: bool = True) -> Iterator[Tuple[str, torch.nn.parameter.Parameter]]\n",
      " |      Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.\n",
      " |      \n",
      " |      Args:\n",
      " |          prefix (str): prefix to prepend to all parameter names.\n",
      " |          recurse (bool): if True, then yields parameters of this module\n",
      " |              and all submodules. Otherwise, yields only parameters that\n",
      " |              are direct members of this module.\n",
      " |          remove_duplicate (bool, optional): whether to remove the duplicated\n",
      " |              parameters in the result. Defaults to True.\n",
      " |      \n",
      " |      Yields:\n",
      " |          (str, Parameter): Tuple containing the name and parameter\n",
      " |      \n",
      " |      Example::\n",
      " |      \n",
      " |          >>> # xdoctest: +SKIP(\"undefined vars\")\n",
      " |          >>> for name, param in self.named_parameters():\n",
      " |          >>>     if name in ['bias']:\n",
      " |          >>>         print(param.size())\n",
      " |  \n",
      " |  parameters(self, recurse: bool = True) -> Iterator[torch.nn.parameter.Parameter]\n",
      " |      Return an iterator over module parameters.\n",
      " |      \n",
      " |      This is typically passed to an optimizer.\n",
      " |      \n",
      " |      Args:\n",
      " |          recurse (bool): if True, then yields parameters of this module\n",
      " |              and all submodules. Otherwise, yields only parameters that\n",
      " |              are direct members of this module.\n",
      " |      \n",
      " |      Yields:\n",
      " |          Parameter: module parameter\n",
      " |      \n",
      " |      Example::\n",
      " |      \n",
      " |          >>> # xdoctest: +SKIP(\"undefined vars\")\n",
      " |          >>> for param in model.parameters():\n",
      " |          >>>     print(type(param), param.size())\n",
      " |          <class 'torch.Tensor'> (20L,)\n",
      " |          <class 'torch.Tensor'> (20L, 1L, 5L, 5L)\n",
      " |  \n",
      " |  register_backward_hook(self, hook: Callable[[ForwardRef('Module'), Union[Tuple[torch.Tensor, ...], torch.Tensor], Union[Tuple[torch.Tensor, ...], torch.Tensor]], Union[NoneType, Tuple[torch.Tensor, ...], torch.Tensor]]) -> torch.utils.hooks.RemovableHandle\n",
      " |      Register a backward hook on the module.\n",
      " |      \n",
      " |      This function is deprecated in favor of :meth:`~torch.nn.Module.register_full_backward_hook` and\n",
      " |      the behavior of this function will change in future versions.\n",
      " |      \n",
      " |      Returns:\n",
      " |          :class:`torch.utils.hooks.RemovableHandle`:\n",
      " |              a handle that can be used to remove the added hook by calling\n",
      " |              ``handle.remove()``\n",
      " |  \n",
      " |  register_buffer(self, name: str, tensor: Optional[torch.Tensor], persistent: bool = True) -> None\n",
      " |      Add a buffer to the module.\n",
      " |      \n",
      " |      This is typically used to register a buffer that should not to be\n",
      " |      considered a model parameter. For example, BatchNorm's ``running_mean``\n",
      " |      is not a parameter, but is part of the module's state. Buffers, by\n",
      " |      default, are persistent and will be saved alongside parameters. This\n",
      " |      behavior can be changed by setting :attr:`persistent` to ``False``. The\n",
      " |      only difference between a persistent buffer and a non-persistent buffer\n",
      " |      is that the latter will not be a part of this module's\n",
      " |      :attr:`state_dict`.\n",
      " |      \n",
      " |      Buffers can be accessed as attributes using given names.\n",
      " |      \n",
      " |      Args:\n",
      " |          name (str): name of the buffer. The buffer can be accessed\n",
      " |              from this module using the given name\n",
      " |          tensor (Tensor or None): buffer to be registered. If ``None``, then operations\n",
      " |              that run on buffers, such as :attr:`cuda`, are ignored. If ``None``,\n",
      " |              the buffer is **not** included in the module's :attr:`state_dict`.\n",
      " |          persistent (bool): whether the buffer is part of this module's\n",
      " |              :attr:`state_dict`.\n",
      " |      \n",
      " |      Example::\n",
      " |      \n",
      " |          >>> # xdoctest: +SKIP(\"undefined vars\")\n",
      " |          >>> self.register_buffer('running_mean', torch.zeros(num_features))\n",
      " |  \n",
      " |  register_forward_hook(self, hook: Union[Callable[[~T, Tuple[Any, ...], Any], Optional[Any]], Callable[[~T, Tuple[Any, ...], Dict[str, Any], Any], Optional[Any]]], *, prepend: bool = False, with_kwargs: bool = False, always_call: bool = False) -> torch.utils.hooks.RemovableHandle\n",
      " |      Register a forward hook on the module.\n",
      " |      \n",
      " |      The hook will be called every time after :func:`forward` has computed an output.\n",
      " |      \n",
      " |      If ``with_kwargs`` is ``False`` or not specified, the input contains only\n",
      " |      the positional arguments given to the module. Keyword arguments won't be\n",
      " |      passed to the hooks and only to the ``forward``. The hook can modify the\n",
      " |      output. It can modify the input inplace but it will not have effect on\n",
      " |      forward since this is called after :func:`forward` is called. The hook\n",
      " |      should have the following signature::\n",
      " |      \n",
      " |          hook(module, args, output) -> None or modified output\n",
      " |      \n",
      " |      If ``with_kwargs`` is ``True``, the forward hook will be passed the\n",
      " |      ``kwargs`` given to the forward function and be expected to return the\n",
      " |      output possibly modified. The hook should have the following signature::\n",
      " |      \n",
      " |          hook(module, args, kwargs, output) -> None or modified output\n",
      " |      \n",
      " |      Args:\n",
      " |          hook (Callable): The user defined hook to be registered.\n",
      " |          prepend (bool): If ``True``, the provided ``hook`` will be fired\n",
      " |              before all existing ``forward`` hooks on this\n",
      " |              :class:`torch.nn.modules.Module`. Otherwise, the provided\n",
      " |              ``hook`` will be fired after all existing ``forward`` hooks on\n",
      " |              this :class:`torch.nn.modules.Module`. Note that global\n",
      " |              ``forward`` hooks registered with\n",
      " |              :func:`register_module_forward_hook` will fire before all hooks\n",
      " |              registered by this method.\n",
      " |              Default: ``False``\n",
      " |          with_kwargs (bool): If ``True``, the ``hook`` will be passed the\n",
      " |              kwargs given to the forward function.\n",
      " |              Default: ``False``\n",
      " |          always_call (bool): If ``True`` the ``hook`` will be run regardless of\n",
      " |              whether an exception is raised while calling the Module.\n",
      " |              Default: ``False``\n",
      " |      \n",
      " |      Returns:\n",
      " |          :class:`torch.utils.hooks.RemovableHandle`:\n",
      " |              a handle that can be used to remove the added hook by calling\n",
      " |              ``handle.remove()``\n",
      " |  \n",
      " |  register_forward_pre_hook(self, hook: Union[Callable[[~T, Tuple[Any, ...]], Optional[Any]], Callable[[~T, Tuple[Any, ...], Dict[str, Any]], Optional[Tuple[Any, Dict[str, Any]]]]], *, prepend: bool = False, with_kwargs: bool = False) -> torch.utils.hooks.RemovableHandle\n",
      " |      Register a forward pre-hook on the module.\n",
      " |      \n",
      " |      The hook will be called every time before :func:`forward` is invoked.\n",
      " |      \n",
      " |      \n",
      " |      If ``with_kwargs`` is false or not specified, the input contains only\n",
      " |      the positional arguments given to the module. Keyword arguments won't be\n",
      " |      passed to the hooks and only to the ``forward``. The hook can modify the\n",
      " |      input. User can either return a tuple or a single modified value in the\n",
      " |      hook. We will wrap the value into a tuple if a single value is returned\n",
      " |      (unless that value is already a tuple). The hook should have the\n",
      " |      following signature::\n",
      " |      \n",
      " |          hook(module, args) -> None or modified input\n",
      " |      \n",
      " |      If ``with_kwargs`` is true, the forward pre-hook will be passed the\n",
      " |      kwargs given to the forward function. And if the hook modifies the\n",
      " |      input, both the args and kwargs should be returned. The hook should have\n",
      " |      the following signature::\n",
      " |      \n",
      " |          hook(module, args, kwargs) -> None or a tuple of modified input and kwargs\n",
      " |      \n",
      " |      Args:\n",
      " |          hook (Callable): The user defined hook to be registered.\n",
      " |          prepend (bool): If true, the provided ``hook`` will be fired before\n",
      " |              all existing ``forward_pre`` hooks on this\n",
      " |              :class:`torch.nn.modules.Module`. Otherwise, the provided\n",
      " |              ``hook`` will be fired after all existing ``forward_pre`` hooks\n",
      " |              on this :class:`torch.nn.modules.Module`. Note that global\n",
      " |              ``forward_pre`` hooks registered with\n",
      " |              :func:`register_module_forward_pre_hook` will fire before all\n",
      " |              hooks registered by this method.\n",
      " |              Default: ``False``\n",
      " |          with_kwargs (bool): If true, the ``hook`` will be passed the kwargs\n",
      " |              given to the forward function.\n",
      " |              Default: ``False``\n",
      " |      \n",
      " |      Returns:\n",
      " |          :class:`torch.utils.hooks.RemovableHandle`:\n",
      " |              a handle that can be used to remove the added hook by calling\n",
      " |              ``handle.remove()``\n",
      " |  \n",
      " |  register_full_backward_hook(self, hook: Callable[[ForwardRef('Module'), Union[Tuple[torch.Tensor, ...], torch.Tensor], Union[Tuple[torch.Tensor, ...], torch.Tensor]], Union[NoneType, Tuple[torch.Tensor, ...], torch.Tensor]], prepend: bool = False) -> torch.utils.hooks.RemovableHandle\n",
      " |      Register a backward hook on the module.\n",
      " |      \n",
      " |      The hook will be called every time the gradients with respect to a module\n",
      " |      are computed, i.e. the hook will execute if and only if the gradients with\n",
      " |      respect to module outputs are computed. The hook should have the following\n",
      " |      signature::\n",
      " |      \n",
      " |          hook(module, grad_input, grad_output) -> tuple(Tensor) or None\n",
      " |      \n",
      " |      The :attr:`grad_input` and :attr:`grad_output` are tuples that contain the gradients\n",
      " |      with respect to the inputs and outputs respectively. The hook should\n",
      " |      not modify its arguments, but it can optionally return a new gradient with\n",
      " |      respect to the input that will be used in place of :attr:`grad_input` in\n",
      " |      subsequent computations. :attr:`grad_input` will only correspond to the inputs given\n",
      " |      as positional arguments and all kwarg arguments are ignored. Entries\n",
      " |      in :attr:`grad_input` and :attr:`grad_output` will be ``None`` for all non-Tensor\n",
      " |      arguments.\n",
      " |      \n",
      " |      For technical reasons, when this hook is applied to a Module, its forward function will\n",
      " |      receive a view of each Tensor passed to the Module. Similarly the caller will receive a view\n",
      " |      of each Tensor returned by the Module's forward function.\n",
      " |      \n",
      " |      .. warning ::\n",
      " |          Modifying inputs or outputs inplace is not allowed when using backward hooks and\n",
      " |          will raise an error.\n",
      " |      \n",
      " |      Args:\n",
      " |          hook (Callable): The user-defined hook to be registered.\n",
      " |          prepend (bool): If true, the provided ``hook`` will be fired before\n",
      " |              all existing ``backward`` hooks on this\n",
      " |              :class:`torch.nn.modules.Module`. Otherwise, the provided\n",
      " |              ``hook`` will be fired after all existing ``backward`` hooks on\n",
      " |              this :class:`torch.nn.modules.Module`. Note that global\n",
      " |              ``backward`` hooks registered with\n",
      " |              :func:`register_module_full_backward_hook` will fire before\n",
      " |              all hooks registered by this method.\n",
      " |      \n",
      " |      Returns:\n",
      " |          :class:`torch.utils.hooks.RemovableHandle`:\n",
      " |              a handle that can be used to remove the added hook by calling\n",
      " |              ``handle.remove()``\n",
      " |  \n",
      " |  register_full_backward_pre_hook(self, hook: Callable[[ForwardRef('Module'), Union[Tuple[torch.Tensor, ...], torch.Tensor]], Union[NoneType, Tuple[torch.Tensor, ...], torch.Tensor]], prepend: bool = False) -> torch.utils.hooks.RemovableHandle\n",
      " |      Register a backward pre-hook on the module.\n",
      " |      \n",
      " |      The hook will be called every time the gradients for the module are computed.\n",
      " |      The hook should have the following signature::\n",
      " |      \n",
      " |          hook(module, grad_output) -> tuple[Tensor] or None\n",
      " |      \n",
      " |      The :attr:`grad_output` is a tuple. The hook should\n",
      " |      not modify its arguments, but it can optionally return a new gradient with\n",
      " |      respect to the output that will be used in place of :attr:`grad_output` in\n",
      " |      subsequent computations. Entries in :attr:`grad_output` will be ``None`` for\n",
      " |      all non-Tensor arguments.\n",
      " |      \n",
      " |      For technical reasons, when this hook is applied to a Module, its forward function will\n",
      " |      receive a view of each Tensor passed to the Module. Similarly the caller will receive a view\n",
      " |      of each Tensor returned by the Module's forward function.\n",
      " |      \n",
      " |      .. warning ::\n",
      " |          Modifying inputs inplace is not allowed when using backward hooks and\n",
      " |          will raise an error.\n",
      " |      \n",
      " |      Args:\n",
      " |          hook (Callable): The user-defined hook to be registered.\n",
      " |          prepend (bool): If true, the provided ``hook`` will be fired before\n",
      " |              all existing ``backward_pre`` hooks on this\n",
      " |              :class:`torch.nn.modules.Module`. Otherwise, the provided\n",
      " |              ``hook`` will be fired after all existing ``backward_pre`` hooks\n",
      " |              on this :class:`torch.nn.modules.Module`. Note that global\n",
      " |              ``backward_pre`` hooks registered with\n",
      " |              :func:`register_module_full_backward_pre_hook` will fire before\n",
      " |              all hooks registered by this method.\n",
      " |      \n",
      " |      Returns:\n",
      " |          :class:`torch.utils.hooks.RemovableHandle`:\n",
      " |              a handle that can be used to remove the added hook by calling\n",
      " |              ``handle.remove()``\n",
      " |  \n",
      " |  register_load_state_dict_post_hook(self, hook)\n",
      " |      Register a post-hook to be run after module's :meth:`~nn.Module.load_state_dict` is called.\n",
      " |      \n",
      " |      It should have the following signature::\n",
      " |          hook(module, incompatible_keys) -> None\n",
      " |      \n",
      " |      The ``module`` argument is the current module that this hook is registered\n",
      " |      on, and the ``incompatible_keys`` argument is a ``NamedTuple`` consisting\n",
      " |      of attributes ``missing_keys`` and ``unexpected_keys``. ``missing_keys``\n",
      " |      is a ``list`` of ``str`` containing the missing keys and\n",
      " |      ``unexpected_keys`` is a ``list`` of ``str`` containing the unexpected keys.\n",
      " |      \n",
      " |      The given incompatible_keys can be modified inplace if needed.\n",
      " |      \n",
      " |      Note that the checks performed when calling :func:`load_state_dict` with\n",
      " |      ``strict=True`` are affected by modifications the hook makes to\n",
      " |      ``missing_keys`` or ``unexpected_keys``, as expected. Additions to either\n",
      " |      set of keys will result in an error being thrown when ``strict=True``, and\n",
      " |      clearing out both missing and unexpected keys will avoid an error.\n",
      " |      \n",
      " |      Returns:\n",
      " |          :class:`torch.utils.hooks.RemovableHandle`:\n",
      " |              a handle that can be used to remove the added hook by calling\n",
      " |              ``handle.remove()``\n",
      " |  \n",
      " |  register_load_state_dict_pre_hook(self, hook)\n",
      " |      Register a pre-hook to be run before module's :meth:`~nn.Module.load_state_dict` is called.\n",
      " |      \n",
      " |      It should have the following signature::\n",
      " |          hook(module, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) -> None  # noqa: B950\n",
      " |      \n",
      " |      Arguments:\n",
      " |          hook (Callable): Callable hook that will be invoked before\n",
      " |              loading the state dict.\n",
      " |  \n",
      " |  register_module(self, name: str, module: Optional[ForwardRef('Module')]) -> None\n",
      " |      Alias for :func:`add_module`.\n",
      " |  \n",
      " |  register_parameter(self, name: str, param: Optional[torch.nn.parameter.Parameter]) -> None\n",
      " |      Add a parameter to the module.\n",
      " |      \n",
      " |      The parameter can be accessed as an attribute using given name.\n",
      " |      \n",
      " |      Args:\n",
      " |          name (str): name of the parameter. The parameter can be accessed\n",
      " |              from this module using the given name\n",
      " |          param (Parameter or None): parameter to be added to the module. If\n",
      " |              ``None``, then operations that run on parameters, such as :attr:`cuda`,\n",
      " |              are ignored. If ``None``, the parameter is **not** included in the\n",
      " |              module's :attr:`state_dict`.\n",
      " |  \n",
      " |  register_state_dict_post_hook(self, hook)\n",
      " |      Register a post-hook for the :meth:`~torch.nn.Module.state_dict` method.\n",
      " |      \n",
      " |      It should have the following signature::\n",
      " |          hook(module, state_dict, prefix, local_metadata) -> None\n",
      " |      \n",
      " |      The registered hooks can modify the ``state_dict`` inplace.\n",
      " |  \n",
      " |  register_state_dict_pre_hook(self, hook)\n",
      " |      Register a pre-hook for the :meth:`~torch.nn.Module.state_dict` method.\n",
      " |      \n",
      " |      It should have the following signature::\n",
      " |          hook(module, prefix, keep_vars) -> None\n",
      " |      \n",
      " |      The registered hooks can be used to perform pre-processing before the ``state_dict``\n",
      " |      call is made.\n",
      " |  \n",
      " |  requires_grad_(self: ~T, requires_grad: bool = True) -> ~T\n",
      " |      Change if autograd should record operations on parameters in this module.\n",
      " |      \n",
      " |      This method sets the parameters' :attr:`requires_grad` attributes\n",
      " |      in-place.\n",
      " |      \n",
      " |      This method is helpful for freezing part of the module for finetuning\n",
      " |      or training parts of a model individually (e.g., GAN training).\n",
      " |      \n",
      " |      See :ref:`locally-disable-grad-doc` for a comparison between\n",
      " |      `.requires_grad_()` and several similar mechanisms that may be confused with it.\n",
      " |      \n",
      " |      Args:\n",
      " |          requires_grad (bool): whether autograd should record operations on\n",
      " |                                parameters in this module. Default: ``True``.\n",
      " |      \n",
      " |      Returns:\n",
      " |          Module: self\n",
      " |  \n",
      " |  set_extra_state(self, state: Any) -> None\n",
      " |      Set extra state contained in the loaded `state_dict`.\n",
      " |      \n",
      " |      This function is called from :func:`load_state_dict` to handle any extra state\n",
      " |      found within the `state_dict`. Implement this function and a corresponding\n",
      " |      :func:`get_extra_state` for your module if you need to store extra state within its\n",
      " |      `state_dict`.\n",
      " |      \n",
      " |      Args:\n",
      " |          state (dict): Extra state from the `state_dict`\n",
      " |  \n",
      " |  set_submodule(self, target: str, module: 'Module') -> None\n",
      " |      Set the submodule given by ``target`` if it exists, otherwise throw an error.\n",
      " |      \n",
      " |      For example, let's say you have an ``nn.Module`` ``A`` that\n",
      " |      looks like this:\n",
      " |      \n",
      " |      .. code-block:: text\n",
      " |      \n",
      " |          A(\n",
      " |              (net_b): Module(\n",
      " |                  (net_c): Module(\n",
      " |                      (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))\n",
      " |                  )\n",
      " |                  (linear): Linear(in_features=100, out_features=200, bias=True)\n",
      " |              )\n",
      " |          )\n",
      " |      \n",
      " |      (The diagram shows an ``nn.Module`` ``A``. ``A`` has a nested\n",
      " |      submodule ``net_b``, which itself has two submodules ``net_c``\n",
      " |      and ``linear``. ``net_c`` then has a submodule ``conv``.)\n",
      " |      \n",
      " |      To overide the ``Conv2d`` with a new submodule ``Linear``, you\n",
      " |      would call\n",
      " |      ``set_submodule(\"net_b.net_c.conv\", nn.Linear(33, 16))``.\n",
      " |      \n",
      " |      Args:\n",
      " |          target: The fully-qualified string name of the submodule\n",
      " |              to look for. (See above example for how to specify a\n",
      " |              fully-qualified string.)\n",
      " |          module: The module to set the submodule to.\n",
      " |      \n",
      " |      Raises:\n",
      " |          ValueError: If the target string is empty\n",
      " |          AttributeError: If the target string references an invalid\n",
      " |              path or resolves to something that is not an\n",
      " |              ``nn.Module``\n",
      " |  \n",
      " |  share_memory(self: ~T) -> ~T\n",
      " |      See :meth:`torch.Tensor.share_memory_`.\n",
      " |  \n",
      " |  state_dict(self, *args, destination=None, prefix='', keep_vars=False)\n",
      " |      Return a dictionary containing references to the whole state of the module.\n",
      " |      \n",
      " |      Both parameters and persistent buffers (e.g. running averages) are\n",
      " |      included. Keys are corresponding parameter and buffer names.\n",
      " |      Parameters and buffers set to ``None`` are not included.\n",
      " |      \n",
      " |      .. note::\n",
      " |          The returned object is a shallow copy. It contains references\n",
      " |          to the module's parameters and buffers.\n",
      " |      \n",
      " |      .. warning::\n",
      " |          Currently ``state_dict()`` also accepts positional arguments for\n",
      " |          ``destination``, ``prefix`` and ``keep_vars`` in order. However,\n",
      " |          this is being deprecated and keyword arguments will be enforced in\n",
      " |          future releases.\n",
      " |      \n",
      " |      .. warning::\n",
      " |          Please avoid the use of argument ``destination`` as it is not\n",
      " |          designed for end-users.\n",
      " |      \n",
      " |      Args:\n",
      " |          destination (dict, optional): If provided, the state of module will\n",
      " |              be updated into the dict and the same object is returned.\n",
      " |              Otherwise, an ``OrderedDict`` will be created and returned.\n",
      " |              Default: ``None``.\n",
      " |          prefix (str, optional): a prefix added to parameter and buffer\n",
      " |              names to compose the keys in state_dict. Default: ``''``.\n",
      " |          keep_vars (bool, optional): by default the :class:`~torch.Tensor` s\n",
      " |              returned in the state dict are detached from autograd. If it's\n",
      " |              set to ``True``, detaching will not be performed.\n",
      " |              Default: ``False``.\n",
      " |      \n",
      " |      Returns:\n",
      " |          dict:\n",
      " |              a dictionary containing a whole state of the module\n",
      " |      \n",
      " |      Example::\n",
      " |      \n",
      " |          >>> # xdoctest: +SKIP(\"undefined vars\")\n",
      " |          >>> module.state_dict().keys()\n",
      " |          ['bias', 'weight']\n",
      " |  \n",
      " |  to(self, *args, **kwargs)\n",
      " |      Move and/or cast the parameters and buffers.\n",
      " |      \n",
      " |      This can be called as\n",
      " |      \n",
      " |      .. function:: to(device=None, dtype=None, non_blocking=False)\n",
      " |         :noindex:\n",
      " |      \n",
      " |      .. function:: to(dtype, non_blocking=False)\n",
      " |         :noindex:\n",
      " |      \n",
      " |      .. function:: to(tensor, non_blocking=False)\n",
      " |         :noindex:\n",
      " |      \n",
      " |      .. function:: to(memory_format=torch.channels_last)\n",
      " |         :noindex:\n",
      " |      \n",
      " |      Its signature is similar to :meth:`torch.Tensor.to`, but only accepts\n",
      " |      floating point or complex :attr:`dtype`\\ s. In addition, this method will\n",
      " |      only cast the floating point or complex parameters and buffers to :attr:`dtype`\n",
      " |      (if given). The integral parameters and buffers will be moved\n",
      " |      :attr:`device`, if that is given, but with dtypes unchanged. When\n",
      " |      :attr:`non_blocking` is set, it tries to convert/move asynchronously\n",
      " |      with respect to the host if possible, e.g., moving CPU Tensors with\n",
      " |      pinned memory to CUDA devices.\n",
      " |      \n",
      " |      See below for examples.\n",
      " |      \n",
      " |      .. note::\n",
      " |          This method modifies the module in-place.\n",
      " |      \n",
      " |      Args:\n",
      " |          device (:class:`torch.device`): the desired device of the parameters\n",
      " |              and buffers in this module\n",
      " |          dtype (:class:`torch.dtype`): the desired floating point or complex dtype of\n",
      " |              the parameters and buffers in this module\n",
      " |          tensor (torch.Tensor): Tensor whose dtype and device are the desired\n",
      " |              dtype and device for all parameters and buffers in this module\n",
      " |          memory_format (:class:`torch.memory_format`): the desired memory\n",
      " |              format for 4D parameters and buffers in this module (keyword\n",
      " |              only argument)\n",
      " |      \n",
      " |      Returns:\n",
      " |          Module: self\n",
      " |      \n",
      " |      Examples::\n",
      " |      \n",
      " |          >>> # xdoctest: +IGNORE_WANT(\"non-deterministic\")\n",
      " |          >>> linear = nn.Linear(2, 2)\n",
      " |          >>> linear.weight\n",
      " |          Parameter containing:\n",
      " |          tensor([[ 0.1913, -0.3420],\n",
      " |                  [-0.5113, -0.2325]])\n",
      " |          >>> linear.to(torch.double)\n",
      " |          Linear(in_features=2, out_features=2, bias=True)\n",
      " |          >>> linear.weight\n",
      " |          Parameter containing:\n",
      " |          tensor([[ 0.1913, -0.3420],\n",
      " |                  [-0.5113, -0.2325]], dtype=torch.float64)\n",
      " |          >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA1)\n",
      " |          >>> gpu1 = torch.device(\"cuda:1\")\n",
      " |          >>> linear.to(gpu1, dtype=torch.half, non_blocking=True)\n",
      " |          Linear(in_features=2, out_features=2, bias=True)\n",
      " |          >>> linear.weight\n",
      " |          Parameter containing:\n",
      " |          tensor([[ 0.1914, -0.3420],\n",
      " |                  [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')\n",
      " |          >>> cpu = torch.device(\"cpu\")\n",
      " |          >>> linear.to(cpu)\n",
      " |          Linear(in_features=2, out_features=2, bias=True)\n",
      " |          >>> linear.weight\n",
      " |          Parameter containing:\n",
      " |          tensor([[ 0.1914, -0.3420],\n",
      " |                  [-0.5112, -0.2324]], dtype=torch.float16)\n",
      " |      \n",
      " |          >>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)\n",
      " |          >>> linear.weight\n",
      " |          Parameter containing:\n",
      " |          tensor([[ 0.3741+0.j,  0.2382+0.j],\n",
      " |                  [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)\n",
      " |          >>> linear(torch.ones(3, 2, dtype=torch.cdouble))\n",
      " |          tensor([[0.6122+0.j, 0.1150+0.j],\n",
      " |                  [0.6122+0.j, 0.1150+0.j],\n",
      " |                  [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)\n",
      " |  \n",
      " |  to_empty(self: ~T, *, device: Union[int, str, torch.device, NoneType], recurse: bool = True) -> ~T\n",
      " |      Move the parameters and buffers to the specified device without copying storage.\n",
      " |      \n",
      " |      Args:\n",
      " |          device (:class:`torch.device`): The desired device of the parameters\n",
      " |              and buffers in this module.\n",
      " |          recurse (bool): Whether parameters and buffers of submodules should\n",
      " |              be recursively moved to the specified device.\n",
      " |      \n",
      " |      Returns:\n",
      " |          Module: self\n",
      " |  \n",
      " |  train(self: ~T, mode: bool = True) -> ~T\n",
      " |      Set the module in training mode.\n",
      " |      \n",
      " |      This has an effect only on certain modules. See the documentation of\n",
      " |      particular modules for details of their behaviors in training/evaluation\n",
      " |      mode, i.e., whether they are affected, e.g. :class:`Dropout`, :class:`BatchNorm`,\n",
      " |      etc.\n",
      " |      \n",
      " |      Args:\n",
      " |          mode (bool): whether to set training mode (``True``) or evaluation\n",
      " |                       mode (``False``). Default: ``True``.\n",
      " |      \n",
      " |      Returns:\n",
      " |          Module: self\n",
      " |  \n",
      " |  type(self: ~T, dst_type: Union[torch.dtype, str]) -> ~T\n",
      " |      Casts all parameters and buffers to :attr:`dst_type`.\n",
      " |      \n",
      " |      .. note::\n",
      " |          This method modifies the module in-place.\n",
      " |      \n",
      " |      Args:\n",
      " |          dst_type (type or string): the desired type\n",
      " |      \n",
      " |      Returns:\n",
      " |          Module: self\n",
      " |  \n",
      " |  xpu(self: ~T, device: Union[int, torch.device, NoneType] = None) -> ~T\n",
      " |      Move all model parameters and buffers to the XPU.\n",
      " |      \n",
      " |      This also makes associated parameters and buffers different objects. So\n",
      " |      it should be called before constructing optimizer if the module will\n",
      " |      live on XPU while being optimized.\n",
      " |      \n",
      " |      .. note::\n",
      " |          This method modifies the module in-place.\n",
      " |      \n",
      " |      Arguments:\n",
      " |          device (int, optional): if specified, all parameters will be\n",
      " |              copied to that device\n",
      " |      \n",
      " |      Returns:\n",
      " |          Module: self\n",
      " |  \n",
      " |  zero_grad(self, set_to_none: bool = True) -> None\n",
      " |      Reset gradients of all model parameters.\n",
      " |      \n",
      " |      See similar function under :class:`torch.optim.Optimizer` for more context.\n",
      " |      \n",
      " |      Args:\n",
      " |          set_to_none (bool): instead of setting to zero, set the grads to None.\n",
      " |              See :meth:`torch.optim.Optimizer.zero_grad` for details.\n",
      " |  \n",
      " |  ----------------------------------------------------------------------\n",
      " |  Data descriptors inherited from torch.nn.modules.module.Module:\n",
      " |  \n",
      " |  __dict__\n",
      " |      dictionary for instance variables\n",
      " |  \n",
      " |  __weakref__\n",
      " |      list of weak references to the object\n",
      " |  \n",
      " |  ----------------------------------------------------------------------\n",
      " |  Data and other attributes inherited from torch.nn.modules.module.Module:\n",
      " |  \n",
      " |  T_destination = ~T_destination\n",
      " |  \n",
      " |  call_super_init = False\n",
      " |  \n",
      " |  dump_patches = False\n",
      "\n"
     ]
    }
   ],
   "source": [
    "help(rnn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "import torch\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([768, 7, 32])\n",
      "torch.Size([32, 7, 24, 1, 32])\n"
     ]
    }
   ],
   "source": [
    "c_in = 1\n",
    "dim_time = 4\n",
    "c_out = 32\n",
    "R = 7\n",
    "P = 24\n",
    "rnn = nn.GRU(c_in * (1 + dim_time), c_in * c_out, num_layers=1, batch_first=True)\n",
    "x = torch.randn(32, R, P, c_in, 1+dim_time).permute(0, 2, 1, 3, 4).reshape(32*P, R, c_in*(1+dim_time))\n",
    "\n",
    "out, h_n = rnn(x)\n",
    "print(out.shape)\n",
    "\n",
    "out = out.reshape(32, P, R, c_in, c_out).permute(0, 2, 1, 3, 4)\n",
    "print(out.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([768, 7, 32])\n",
      "torch.Size([32, 7, 24, 1, 32])\n"
     ]
    }
   ],
   "source": [
    "c_in = 1\n",
    "dim_time = 4\n",
    "c_out = 32\n",
    "R = 7\n",
    "P = 24\n",
    "rnn = nn.LSTM(c_in * (1 + dim_time), c_in * c_out, num_layers=1, batch_first=True)\n",
    "x = torch.randn(32, R, P, c_in, 1+dim_time).permute(0, 2, 1, 3, 4).reshape(32*P, R, c_in*(1+dim_time))\n",
    "\n",
    "out, h_n = rnn(x)\n",
    "print(out.shape)\n",
    "\n",
    "out = out.reshape(32, P, R, c_in, c_out).permute(0, 2, 1, 3, 4)\n",
    "print(out.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([768, 7, 32])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "torch.Size([32, 7, 24, 1, 32])"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "c_in = 1\n",
    "dim_time = 4\n",
    "c_out = 32\n",
    "R = 7\n",
    "P = 24\n",
    "mlp = nn.Sequential(\n",
    "                nn.Linear(c_in * (1 + dim_time), c_in * c_out),\n",
    "                nn.ReLU()\n",
    "            )\n",
    "x = torch.randn(32, R, P, c_in, 1+dim_time).permute(0, 2, 1, 3, 4).reshape(32*P, R, c_in*(1+dim_time))\n",
    "\n",
    "out = mlp(x)\n",
    "print(out.shape)\n",
    "\n",
    "out = out.reshape(32, P, R, c_in, c_out).permute(0, 2, 1, 3, 4)\n",
    "(out.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([768, 7, 5])\n",
      "torch.Size([32, 7, 24, 1, 32])\n"
     ]
    }
   ],
   "source": [
    "c_in = 1\n",
    "dim_time = 4\n",
    "c_out = 32\n",
    "R = 7\n",
    "P = 24\n",
    "attn = nn.MultiheadAttention(c_in * (1 + dim_time), 1, batch_first=True)\n",
    "linear = nn.Linear(c_in * (1 + dim_time), c_in * c_out)\n",
    "dropout = nn.Dropout(0.5)\n",
    "x = torch.randn(32, R, P, c_in, 1+dim_time).permute(0, 2, 1, 3, 4).reshape(32*P, R, c_in*(1+dim_time))\n",
    "\n",
    "out, _ = attn(x,x,x)\n",
    "print(out.shape)\n",
    "\n",
    "out = dropout(linear(out))\n",
    "out = out.reshape(32, P, R, c_in, c_out).permute(0, 2, 1, 3, 4)\n",
    "print(out.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "jjk_for_torch",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
